#!/usr/bin/env python
import os
import socket
import sys
from pathlib import Path

import numpy as np
import setproctitle
import torch

import wandb
from onpolicy.config import get_config
from onpolicy.envs.env_wrappers import DummyVecEnv, SubprocVecEnv
from onpolicy.envs.mpe.MPE_env import MPEEnv


def make_eval_env(all_args):
    def get_env_fn(rank):
        def init_env():
            if all_args.env_name == "MPE":
                env = MPEEnv(all_args)
            else:
                print("Can not support the " + all_args.env_name + "environment.")
                raise NotImplementedError
            env.seed(all_args.seed + rank * 1000)
            return env

        return init_env

    if all_args.n_rollout_threads == 1:
        return DummyVecEnv([get_env_fn(0)])
    else:
        return SubprocVecEnv([get_env_fn(i) for i in range(all_args.n_rollout_threads)])


def parse_args(args, parser):
    parser.add_argument(
        "--scenario_name",
        type=str,
        default="simple_spread",
        help="Which scenario to run on",
    )
    parser.add_argument("--num_landmarks", type=int, default=3)
    parser.add_argument("--num_agents", type=int, default=2, help="number of players")

    all_args = parser.parse_known_args(args)[0]

    return all_args


def main(args):
    parser = get_config()
    all_args = parse_args(args, parser)

    if all_args.algorithm_name == "rmappo":
        print("u are choosing to use rmappo, we set use_recurrent_policy to be True")
        all_args.use_recurrent_policy = True
        all_args.use_naive_recurrent_policy = False
    elif all_args.algorithm_name == "mappo":
        print(
            "u are choosing to use mappo, we set use_recurrent_policy & use_naive_recurrent_policy to be False"
        )
        all_args.use_recurrent_policy = False
        all_args.use_naive_recurrent_policy = False
    elif all_args.algorithm_name == "ippo":
        print("u are choosing to use ippo, we set use_centralized_V to be False.")
        all_args.use_centralized_V = False
    elif all_args.algorithm_name == "mat":
        print("u are choosing to use mat.")
    else:
        raise NotImplementedError

    assert (
        (
            all_args.share_policy == True
            and all_args.scenario_name == "simple_speaker_listener"
        )
        == False
    ), "The simple_speaker_listener scenario can not use shared policy. Please check the config.py."

    # assert all_args.use_render, "u need to set use_render be True"
    assert not (
        all_args.model_dir == None or all_args.model_dir == ""
    ), "set model_dir first"
    # assert all_args.n_rollout_threads == 1, "only support to use 1 env to render."

    # cuda
    if all_args.cuda and torch.cuda.is_available():
        print("choose to use gpu...")
        device = torch.device("cuda:0")
        torch.set_num_threads(all_args.n_training_threads)
        if all_args.cuda_deterministic:
            torch.backends.cudnn.benchmark = False
            torch.backends.cudnn.deterministic = True
    else:
        print("choose to use cpu...")
        device = torch.device("cpu")
        torch.set_num_threads(all_args.n_training_threads)

    # run dir
    run_dir = (
        Path(os.path.split(os.path.dirname(os.path.abspath(__file__)))[0] + "/results")
        / all_args.env_name
        / all_args.scenario_name
        / all_args.algorithm_name
        / all_args.experiment_name
    )
    # wandb

    if all_args.use_wandb:
        run = wandb.init(
            config=all_args,
            project=all_args.wandb_env_name,
            # entity=all_args.user_name,
            notes=socket.gethostname(),
            name=str(all_args.algorithm_name)
            + "_"
            + str(all_args.ele_radius)
            + "_"
            + str(all_args.ele_time),
            group=all_args.scenario_name,
            dir=str(run_dir),
            job_type="evaluation",
            reinit=True,
        )
    else:
        if not run_dir.exists():
            curr_run = (
                "run-1"
                + "-train-"
                + str(all_args.train_ele_radius)
                + "-"
                + str(all_args.train_ele_time)
                + "-eval-"
                + str(all_args.ele_radius)
                + "-"
                + str(all_args.ele_time)
            )
        else:
            exst_run_nums = [
                int(str(folder.name).split("-")[1])
                for folder in run_dir.iterdir()
                if str(folder.name).startswith("run")
            ]
            if len(exst_run_nums) == 0:
                curr_run = (
                    "run-1"
                    + "-train-"
                    + str(all_args.train_ele_radius)
                    + "-"
                    + str(all_args.train_ele_time)
                    + "-eval-"
                    + str(all_args.ele_radius)
                    + "-"
                    + str(all_args.ele_time)
                )
            else:
                curr_run = (
                    "run-%i" % (max(exst_run_nums) + 1)
                    + "-train-"
                    + str(all_args.train_ele_radius)
                    + "-"
                    + str(all_args.train_ele_time)
                    + "-eval-"
                    + str(all_args.ele_radius)
                    + "-"
                    + str(all_args.ele_time)
                )
        run_dir = run_dir / curr_run
        if not run_dir.exists():
            os.makedirs(str(run_dir))

    setproctitle.setproctitle(
        str(all_args.algorithm_name)
        + "-"
        + str(all_args.env_name)
        + "-"
        + str(all_args.experiment_name)
        + "@"
        + str(all_args.user_name)
    )

    # seed
    torch.manual_seed(all_args.seed)
    torch.cuda.manual_seed_all(all_args.seed)
    np.random.seed(all_args.seed)

    # env init
    envs = make_eval_env(all_args)
    eval_envs = None
    num_agents = all_args.num_agents

    config = {
        "all_args": all_args,
        "envs": envs,
        "eval_envs": eval_envs,
        "num_agents": num_agents,
        "device": device,
        "run_dir": run_dir,
    }

    # run experiments
    if all_args.share_policy:
        from onpolicy.runner.shared.mpe_runner import MPERunner as Runner
    else:
        from onpolicy.runner.separated.mpe_runner import MPERunner as Runner

    runner = Runner(config)
    runner.eval_my()

    # post process
    envs.close()
    if all_args.use_wandb:
        run.finish()


if __name__ == "__main__":
    main(sys.argv[1:])
